# SPDX-FileCopyrightText: 2025 LakeSoul Contributors
#
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations
from collections import deque
import sys
from types import ModuleType
from typing import Any, Generic

from annotated_types import T
from e2e.models.dag import DAG


class ContextStackMeta(type):
    _context: deque

    def __new__(
        cls, name, bases, namespace, share_parent_context: bool = False, **kwargs: Any
    ):
        if not share_parent_context:
            namespace["_context"] = deque()
        new_cls = super().__new__(cls, name, bases, namespace, **kwargs)

        return new_cls

    @property
    def active(self) -> bool:
        return bool(self._context)


class ContextStack(Generic[T], metaclass=ContextStackMeta):
    _context: deque[T]

    @classmethod
    def push(cls, obj: T):
        cls._context.appendleft(obj)

    @classmethod
    def pop(cls) -> T | None:
        return cls._context.popleft()

    @classmethod
    def get_current(cls) -> T | None:
        try:
            return cls._context[0]
        except IndexError:
            return None


class DagContext(ContextStack[DAG]):
    """
    DAG context is used to keep the current DAG when DAG is used as ContextManager.

    You can use DAG as context:

    .. code-block:: python

        with DAG(
            dag_id="example_dag",
            default_args=default_args,
            schedule="0 0 * * *",
            dagrun_timeout=timedelta(minutes=60),
        ) as dag:
            ...

    If you do this the context stores the DAG and whenever new task is created, it will use
    such stored DAG as the parent DAG.

    """

    autoregistered_dags: set[tuple[DAG, ModuleType]] = set()
    current_autoregister_module_name: str | None = None

    @classmethod
    def pop(cls) -> DAG | None:
        dag = super().pop()
        # In a few cases around serialization we explicitly push None in to the stack
        if (
            cls.current_autoregister_module_name is not None
            and dag
            and getattr(dag, "auto_register", True)
        ):
            mod = sys.modules[cls.current_autoregister_module_name]
            cls.autoregistered_dags.add((dag, mod))
        return dag

    @classmethod
    def get_current_dag(cls) -> DAG | None:
        return cls.get_current()
